package com.tang.state;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.scala.typeutils.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

/**
 * 在map算子中计算数据的个数
 * 算子状态中，list与unionList，并行度改变，怎么重新分配状态
 * list状态：轮询均分给新的并行子任务
 * unionList状态：原先的多个子任务的状态，合并成完整的一份。广播给新的并行子任务，每个任务一份完整的
 *
 * @author tang
 * @since 2023/7/11 20:41
 */
public class OperatorListStateDemo {

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(new Configuration());
        env.setParallelism(2);

        env.socketTextStream("192.168.70.141", 7777)
                .map(new MyCountMapFunction())
                .print();

        env.execute();
    }

    /**
     * 1.实现CheckpointedFunction接口
     * 不是太能理解，继承这个干了啥？
     */
    private static class MyCountMapFunction implements MapFunction<String, Long>, CheckpointedFunction {

        /**
         * count变量在程序挂掉等情况下会丢失数据
         * Checkpoint会保存快照到算子状态中，在程序重启时调用initializeState方法，我们可以在那里重新给count赋值
         */
        private Long count = 0L;
        private ListState<Long> listState;

        @Override
        public Long map(String value) {
            return ++count;
        }

        /**
         * 对状态做快照
         * 2. 本地变量持久化：将 本地变量 拷贝到算子状态中，【开启checkpoint才会调用后面章节介绍】
         *
         * @param context the context for drawing a snapshot of the operator
         * @throws Exception -
         */
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            System.out.println("SnapshotState...");
            // 2.1清空算子状态，避免被其他的干扰
            listState.clear();
            // 2.2 将 本地变量添加到算子状态中
            listState.add(count);
        }

        /**
         * 初始化状态，用在程序启动、恢复时（每个子任务都会执行一次）
         * 3.初始化本地变量：程序启动和恢复时，从状态中 把数据添加到 本地变量，每个子任务调用一次。
         *
         * @param context the context for initializing the operator
         * @throws Exception -
         */
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            System.out.println("InitializeState...");
            // 3.1 从上下文初始化算子状态
            listState = context.getOperatorStateStore()
                    //.getListState(new ListStateDescriptor<>("state", Types.LONG()))
                    .getUnionListState(new ListStateDescriptor<>("state", Types.LONG()));

            // 3.2 从算子状态中把数据拷贝到本地变量
            if (context.isRestored()) { // 是否恢复成功
                for (Long c : listState.get()) {
                    count += c;
                }
            }
        }
    }

}
